{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "00b35eb1",
   "metadata": {},
   "source": [
    "# Zep Knowledge Graph + LLM Demo\n",
    "\n",
    "This notebook demonstrates **Retrieval-Augmented Generation (RAG)** using:\n",
    "- **Zep Cloud** for knowledge graph storage and semantic search\n",
    "- **OpenRouter** + **Minimax M2** for generating intelligent, coherent responses\n",
    "\n",
    "## What You'll Learn\n",
    "\n",
    "1. **Ingest data** into Zep's Knowledge Graph as episodes\n",
    "2. **Retrieve information** using semantic search\n",
    "3. **Generate natural responses** using an LLM with retrieved context\n",
    "\n",
    "This creates AI assistants that are both **knowledgeable** (grounded in your data) and **conversational** (natural to interact with)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dc7182c",
   "metadata": {},
   "source": [
    "## Setup and Installation\n",
    "\n",
    "First, we'll import the required libraries and set up our clients for both Zep and OpenRouter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4490d410",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from zep_cloud.client import Zep\n",
    "import time\n",
    "\n",
    "# Import OpenAI client for OpenRouter (Minimax M2)\n",
    "from openai import OpenAI\n",
    "\n",
    "# Load environment variables from .env file\n",
    "# This should contain your ZEP_API_KEY and OPENROUTER_API_KEY\n",
    "load_dotenv()\n",
    "\n",
    "# Get API keys from environment\n",
    "ZEP_API_KEY = os.getenv(\"ZEP_API_KEY\")\n",
    "OPENROUTER_API_KEY = os.getenv(\"OPENROUTER_API_KEY\")\n",
    "\n",
    "if not ZEP_API_KEY:\n",
    "    raise ValueError(\"ZEP_API_KEY not found in environment variables. Please add it to your .env file\")\n",
    "\n",
    "if not OPENROUTER_API_KEY:\n",
    "    raise ValueError(\"OPENROUTER_API_KEY not found in environment variables. Please add it to your .env file\")\n",
    "\n",
    "print(\"Environment loaded successfully\")\n",
    "print(f\"Zep API Key found: {ZEP_API_KEY[:5]}...\")\n",
    "print(f\"OpenRouter API Key found: {OPENROUTER_API_KEY[:5]}...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e08328e7",
   "metadata": {},
   "source": [
    "## Initialize Zep Client\n",
    "\n",
    "We create a Zep client instance that will handle all our interactions with the Zep API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff5f9951",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the Zep client\n",
    "# This is the main interface for all Zep operations\n",
    "zep = Zep(api_key=ZEP_API_KEY)\n",
    "\n",
    "# Initialize OpenRouter client for Minimax M2 LLM\n",
    "# We'll use this to generate intelligent responses based on retrieved knowledge\n",
    "openrouter_client = OpenAI(\n",
    "    base_url=\"https://openrouter.ai/api/v1\",\n",
    "    api_key=OPENROUTER_API_KEY\n",
    ")\n",
    "\n",
    "print(\"Zep client initialized successfully\")\n",
    "print(\"OpenRouter client initialized (Minimax M2 ready)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "630a90a5",
   "metadata": {},
   "source": [
    "## Step 1: Create a Demo User\n",
    "\n",
    "In Zep, all data is associated with a **User**. Let's create a demo user for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2accef90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a unique user ID for this demo\n",
    "DEMO_USER_ID = \"demo-user-coffee-shop\"\n",
    "\n",
    "# Try to create the user\n",
    "# If the user already exists, Zep will return an error, which we'll catch\n",
    "try:\n",
    "    user = zep.user.add(\n",
    "        user_id=DEMO_USER_ID,\n",
    "        email=\"demo@coffeeshop.local\",\n",
    "        first_name=\"Demo\",\n",
    "        last_name=\"User\"\n",
    "    )\n",
    "    print(f\"User created: {user.user_id}\")\n",
    "    print(f\"  Name: {user.first_name} {user.last_name}\")\n",
    "    print(f\"  Email: {user.email}\")\n",
    "except Exception as e:\n",
    "    # User likely already exists\n",
    "    if \"already exists\" in str(e).lower() or \"409\" in str(e):\n",
    "        print(f\"User '{DEMO_USER_ID}' already exists (this is fine!)\")\n",
    "    else:\n",
    "        print(f\"Error creating user: {e}\")\n",
    "        print(\"Continuing anyway...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac27c422",
   "metadata": {},
   "source": [
    "## Step 2: Prepare Demo Data\n",
    "\n",
    "Let's create some sample text data about a fictional coffee shop. We'll store each piece of information as a separate **episode** in the knowledge graph.\n",
    "\n",
    "In a real application, episodes could be:\n",
    "- Chunks of documentation\n",
    "- Customer interaction logs\n",
    "- Product descriptions\n",
    "- Research notes\n",
    "- Any unstructured text data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "340b5517",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define our demo episodes about a coffee shop\n",
    "# Each episode is a paragraph of information that will be ingested into the knowledge graph\n",
    "demo_episodes = [\n",
    "    {\n",
    "        \"title\": \"Company Overview\",\n",
    "        \"content\": \"\"\"Bean & Brew is a specialty coffee shop located in downtown Seattle, Washington. \n",
    "        Founded in 2018 by Maria Chen and David Rodriguez, the shop has become a local favorite for \n",
    "        artisanal coffee and pastries. The shop employs 12 staff members and serves approximately \n",
    "        300 customers daily. Bean & Brew is known for its sustainable sourcing practices and \n",
    "        community engagement initiatives.\"\"\"\n",
    "    },\n",
    "    {\n",
    "        \"title\": \"Products and Services\",\n",
    "        \"content\": \"\"\"Bean & Brew offers a wide range of beverages including espresso drinks, pour-over coffee, \n",
    "        cold brew, and specialty teas. The shop sources its coffee beans from small farms in Colombia, \n",
    "        Ethiopia, and Guatemala. In addition to beverages, they serve fresh-baked pastries, sandwiches, \n",
    "        and salads made daily by their in-house baker, Sophie Martin. The shop also sells whole bean \n",
    "        coffee and brewing equipment for home enthusiasts.\"\"\"\n",
    "    },\n",
    "    {\n",
    "        \"title\": \"Operating Hours and Location\",\n",
    "        \"content\": \"\"\"The coffee shop is located at 456 Pine Street, Seattle, WA 98101, right next to \n",
    "        Pioneer Square. Operating hours are Monday through Friday from 6:30 AM to 7:00 PM, Saturday \n",
    "        from 7:00 AM to 8:00 PM, and Sunday from 8:00 AM to 6:00 PM. The shop features free WiFi, \n",
    "        plenty of seating including outdoor patio space, and is wheelchair accessible. Parking is \n",
    "        available in the adjacent public parking garage.\"\"\"\n",
    "    },\n",
    "    {\n",
    "        \"title\": \"Customer Loyalty Program\",\n",
    "        \"content\": \"\"\"Bean & Brew launched its loyalty program, 'Bean Rewards', in January 2020. Members earn \n",
    "        1 point for every dollar spent, and 100 points can be redeemed for a free beverage of any size. \n",
    "        The program has over 2,500 active members. Premium members, who pay an annual fee of $50, receive \n",
    "        20% off all purchases and exclusive access to new menu items. The program is managed through a \n",
    "        mobile app available on iOS and Android.\"\"\"\n",
    "    },\n",
    "    {\n",
    "        \"title\": \"Community Events\",\n",
    "        \"content\": \"\"\"Bean & Brew hosts various community events throughout the year. Every first Saturday, \n",
    "        they organize 'Coffee Cupping Sessions' where customers can learn about different coffee origins \n",
    "        and brewing methods. The shop partners with local artists to display their work on the walls, \n",
    "        with exhibits changing monthly. During summer months, they host 'Music on the Patio' featuring \n",
    "        local musicians every Friday evening from 6 PM to 8 PM. All events are free and open to the public.\"\"\"\n",
    "    }\n",
    "]\n",
    "\n",
    "print(f\"Prepared {len(demo_episodes)} episodes for ingestion\")\n",
    "print(\"\\nEpisode titles:\")\n",
    "for i, episode in enumerate(demo_episodes, 1):\n",
    "    print(f\"  {i}. {episode['title']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91c091cc",
   "metadata": {},
   "source": [
    "## Step 3: Ingest Episodes into Knowledge Graph\n",
    "\n",
    "Now we'll add each episode to the user's knowledge graph. Zep will automatically:\n",
    "- Extract entities (people, places, organizations, dates, etc.)\n",
    "- Identify facts and relationships\n",
    "- Build a searchable knowledge graph\n",
    "- Enable semantic search over the content\n",
    "\n",
    "Each call to `zep.graph.add()` creates an **episode** and returns a unique identifier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faa93144",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store episode UUIDs for later reference\n",
    "episode_uuids = []\n",
    "\n",
    "print(\"Starting ingestion...\\n\")\n",
    "\n",
    "# Iterate through each episode and add it to the knowledge graph\n",
    "for i, episode in enumerate(demo_episodes, 1):\n",
    "    print(f\"[{i}/{len(demo_episodes)}] Ingesting: {episode['title']}\")\n",
    "    \n",
    "    # Format the content with a header for better context\n",
    "    formatted_content = f\"# {episode['title']}\\n\\n{episode['content']}\"\n",
    "    \n",
    "    try:\n",
    "        # Add the episode to the user's knowledge graph\n",
    "        # - user_id: The user this data belongs to\n",
    "        # - type: The format of the data ('text', 'json', or 'message')\n",
    "        # - data: The actual content to ingest\n",
    "        result = zep.graph.add(\n",
    "            user_id=DEMO_USER_ID,\n",
    "            type=\"text\",\n",
    "            data=formatted_content\n",
    "        )\n",
    "        \n",
    "        # Store the episode UUID for reference\n",
    "        episode_uuids.append(result.uuid_)\n",
    "        \n",
    "        print(f\"  Episode created: {result.uuid_}\")\n",
    "        print(f\"  Content length: {len(formatted_content)} characters\\n\")\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"  Error ingesting episode: {e}\\n\")\n",
    "    \n",
    "    # Small delay to avoid rate limiting\n",
    "    time.sleep(0.5)\n",
    "\n",
    "print(f\"\\n{'='*60}\")\n",
    "print(f\"Successfully ingested {len(episode_uuids)} episodes\")\n",
    "print(f\"{'='*60}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee30bf2f",
   "metadata": {},
   "source": [
    "## Step 4: Wait for Graph Processing\n",
    "\n",
    "After ingestion, Zep processes the data in the background to build the knowledge graph. This typically takes a few seconds to a few minutes depending on the amount of data.\n",
    "\n",
    "The graph extraction involves:\n",
    "- Natural language understanding\n",
    "- Entity extraction\n",
    "- Building semantic embeddings for search\n",
    "- etc.\n",
    "\n",
    "Check the status of episodes on the Zep dashboard to see when processing is complete."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9df176bf",
   "metadata": {},
   "source": [
    "## Step 5: Search the Knowledge Graph\n",
    "\n",
    "Now that our data is ingested and processed, we can search the knowledge graph!\n",
    "\n",
    "The `zep.graph.search()` method performs **semantic search**, meaning it understands the meaning of your query, not just keyword matching.\n",
    "\n",
    "Let's try different queries to see what Zep has learned!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f892dace",
   "metadata": {},
   "source": [
    "### Ask Questions with AI-Generated Responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d760829",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ask_with_knowledge(query: str, limit: int = 3) -> str:\n",
    "    \"\"\"\n",
    "    Simple function to query the knowledge graph and use Minimax M2 to generate a coherent response.\n",
    "    \n",
    "    Args:\n",
    "        query: The user's question (natural language)\n",
    "        limit: Maximum number of results to retrieve\n",
    "        \n",
    "    Returns:\n",
    "        A coherent response generated by the LLM based on retrieved knowledge\n",
    "    \"\"\"\n",
    "    print(f\"🔍 Query: '{query}'\")\n",
    "    print(f\"   Retrieving from Zep (limit: {limit})...\")\n",
    "    \n",
    "    try:\n",
    "        # Step 1: Search the knowledge graph\n",
    "        results = zep.graph.search(\n",
    "            user_id=DEMO_USER_ID,\n",
    "            query=query,\n",
    "            scope=\"episodes\",\n",
    "            limit=limit\n",
    "        )\n",
    "        \n",
    "        # Step 2: Extract relevant context from search results\n",
    "        context_parts = []\n",
    "                    \n",
    "        if results.episodes:\n",
    "            print(f\"Found {len(results.episodes)} episodes\")\n",
    "            for episode in results.episodes:\n",
    "                # Include the full episode content for better context\n",
    "                context_parts.append(f\"Document: {episode.content}\")\n",
    "        \n",
    "        if not context_parts:\n",
    "            return \"I couldn't find any relevant information in the knowledge base to answer your question.\"\n",
    "        \n",
    "        # Step 3: Build context for the LLM\n",
    "        retrieved_context = \"\\n\\n\".join(context_parts)\n",
    "        \n",
    "        # Step 4: Create a prompt that combines the query with retrieved context\n",
    "        system_prompt = \"\"\"You are a helpful AI assistant with access to a knowledge base about a coffee shop.\n",
    "                            Your role is to answer questions accurately based on the provided context.\n",
    "\n",
    "                            Guidelines:\n",
    "                            - Answer naturally and conversationally\n",
    "                            - Only use information from the provided context\n",
    "                            - If the context doesn't contain enough information, say so honestly\n",
    "                            - Be concise but informative\n",
    "                            - Format your response in a friendly, helpful way\"\"\"\n",
    "\n",
    "        user_message = f\"\"\"Context from knowledge base:\n",
    "                        {retrieved_context}\n",
    "\n",
    "                        User question: {query}\n",
    "\n",
    "                        Please provide a clear, helpful answer based on the context above.\"\"\"\n",
    "\n",
    "        # Step 5: Call Minimax M2 via OpenRouter to generate response\n",
    "        print(\"Generating response with Minimax M2...\")\n",
    "        \n",
    "        response = openrouter_client.chat.completions.create(\n",
    "            model=\"minimax/minimax-m2\",\n",
    "            messages=[\n",
    "                {\"role\": \"system\", \"content\": system_prompt},\n",
    "                {\"role\": \"user\", \"content\": user_message}\n",
    "            ],\n",
    "            temperature=0.7,\n",
    "            max_tokens=200\n",
    "        )\n",
    "        \n",
    "        # Extract the generated response\n",
    "        if response.choices and response.choices[0].message.content:\n",
    "            answer = response.choices[0].message.content\n",
    "            print(\"Response generated!\\n\")\n",
    "            return answer\n",
    "        else:\n",
    "            return \"Sorry, I couldn't generate a response.\"\n",
    "            \n",
    "    except Exception as e:\n",
    "        print(f\"Error: {e}\")\n",
    "        return f\"An error occurred while processing your question: {str(e)}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f44c642",
   "metadata": {},
   "source": [
    "### Example 1: Ask about the founders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cfdc1aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ask about founders using LLM-powered response\n",
    "response = ask_with_knowledge(\n",
    "    query=\"Who founded Bean & Brew and when was it started?\",\n",
    "    limit=3\n",
    ")\n",
    "\n",
    "print(\"=\" * 70)\n",
    "print(\"AI RESPONSE:\")\n",
    "print(\"=\" * 70)\n",
    "print(response)\n",
    "print(\"=\" * 70)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "490b17b2",
   "metadata": {},
   "source": [
    "### Example 2: Ask about location and hours"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4510935",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get a comprehensive answer about location and hours\n",
    "response = ask_with_knowledge(\n",
    "    query=\"Where is Bean & Brew located and what are their hours?\",\n",
    "    limit=3\n",
    ")\n",
    "\n",
    "print(\"=\" * 70)\n",
    "print(\"AI RESPONSE:\")\n",
    "print(\"=\" * 70)\n",
    "print(response)\n",
    "print(\"=\" * 70)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "168eeac7",
   "metadata": {},
   "source": [
    "### Example 3: Ask about products and menu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a033981",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ask about what's on the menu\n",
    "response = ask_with_knowledge(\n",
    "    query=\"What kinds of drinks and food can I get at Bean & Brew?\",\n",
    "    limit=3\n",
    ")\n",
    "\n",
    "print(\"=\" * 70)\n",
    "print(\"AI RESPONSE:\")\n",
    "print(\"=\" * 70)\n",
    "print(response)\n",
    "print(\"=\" * 70)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6bff0b9",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this notebook, we demonstrated:\n",
    "\n",
    "1. **Creating a Zep user** - Every piece of data in Zep belongs to a user\n",
    "\n",
    "2. **Ingesting episodes** - We added 5 text episodes about a coffee shop to the knowledge graph\n",
    "\n",
    "3. **RAG Pattern (Retrieval-Augmented Generation)** - The key workflow:\n",
    "   - **Retrieve**: Search Zep's knowledge graph for relevant information\n",
    "   - **Augment**: Add retrieved context to the LLM prompt\n",
    "   - **Generate**: Use Minimax M2 to create coherent, natural responses\n",
    "\n",
    "## Technical Stack\n",
    "\n",
    "- **Zep Cloud** - Knowledge graph storage and semantic search\n",
    "- **OpenRouter** - API gateway for accessing various LLMs\n",
    "- **Minimax M2** - Advanced language model for generating responses\n",
    "- **RAG Pattern** - Retrieval-Augmented Generation for accurate, context-aware AI\n",
    "\n",
    "## Resources\n",
    "\n",
    "- [Zep Documentation](https://help.getzep.com/)\n",
    "- [Zep Python SDK](https://github.com/getzep/zep-python)\n",
    "- [OpenRouter API](https://openrouter.ai/)\n",
    "- [Minimax M2 Model](https://openrouter.ai/minimax/minimax-m2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5551dd03",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai-avatar-demo (3.12.10)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
